{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 4: Implementing Self Attention GAN (SAGAN)\n",
    "\n",
    "**Author** - [Avik Pal](https://avik-pal.github.io)\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/torchgan/torchgan/blob/master/tutorials/Tutorial%204.%20Self%20Attention%20GAN.ipynb)\n",
    "\n",
    "In this tutorial we shall be implementing **Self Attention GAN** (SAGAN), which makes use of the novel **Self Attention** layer as proposed in [\n",
    "Self-Attention Generative Adversarial Networks by Zhang et. al.](https://arxiv.org/abs/1805.08318). **Self Attention**, is incorporated in both the **Generator** as well as the **Discriminator**, leading to *state-of-the-art* (at the time when this tutorial was being written) results on ImageNet.\n",
    "\n",
    "The tutorial helps you with the following:\n",
    "\n",
    "1. Understanding how to use the **torchgan.layers API**.\n",
    "2. Creating custom Generator and Discriminator architectures that fit in seamlessly with **TorchGAN'**  by simply extending the base **torchgan.models.Generator** and **torchgan.models.Discriminator** classes\n",
    "3. Visualizing the samples generated on the **CelebA Dataset**. We stick to such a simple dataset for illustration purposes and convergence time issues"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial assumes that your system has **PyTorch** and **TorchGAN** installed properly. If not, the following code block will try to install the **latest tagged version** of TorchGAN. If you need to use some other version head over to the installation instructions on the [official documentation website](https://torchgan.readthedocs.io/en/latest/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import torchgan\n",
    "\n",
    "    print(f\"Existing TorchGAN {torchgan.__version__} installation found\")\n",
    "except ImportError:\n",
    "    import subprocess\n",
    "    import sys\n",
    "\n",
    "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"torchgan\"])\n",
    "    import torchgan\n",
    "\n",
    "    print(f\"Installed TorchGAN {torchgan.__version__}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# General Imports\n",
    "import os\n",
    "import random\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# Pytorch and Torchvision Imports\n",
    "import torch\n",
    "import torchvision\n",
    "from torch.optim import Adam\n",
    "import torch.nn as nn\n",
    "import torch.utils.data as data\n",
    "import torchvision.datasets as dsets\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.utils as vutils\n",
    "\n",
    "# Torchgan Imports\n",
    "import torchgan\n",
    "from torchgan.layers import SpectralNorm2d, SelfAttention2d\n",
    "from torchgan.models import Generator, Discriminator\n",
    "from torchgan.losses import (\n",
    "    WassersteinGeneratorLoss,\n",
    "    WassersteinDiscriminatorLoss,\n",
    "    WassersteinGradientPenalty,\n",
    ")\n",
    "from torchgan.trainer import Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set random seed for reproducibility\n",
    "manualSeed = 144\n",
    "random.seed(manualSeed)\n",
    "torch.manual_seed(manualSeed)\n",
    "print(\"Random Seed: \", manualSeed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LOADING THE DATASET\n",
    "\n",
    "We will try to download the **CelebA Data** in the required format. In case the setup fails for any reason use the following instructions to setup the data (and [open an issue](https://github.com/torchgan/torchgan/issues/new?assignees=&labels=bug&template=bug_report.md&title=)).\n",
    "\n",
    "<br>\n",
    "<details><summary><b>Click to view Data Setup Instructions</b></summary>\n",
    "\n",
    "Before starting this tutorial, download the [CelebA dataset](http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html). Extract the file named **img_align_celeba.zip**. The directory structure should be:\n",
    "\n",
    "```\n",
    "--> path/to/data\n",
    "    --> img_align_celeba\n",
    "        --> 188242.jpg\n",
    "        --> 173822.jpg\n",
    "        --> 284702.jpg\n",
    "        --> 537394.jpg\n",
    "           ...\n",
    "```\n",
    "\n",
    "</details>\n",
    "\n",
    "We make the following transforms before feeding the **CelebA Dataset** into the networks\n",
    "\n",
    "1. By convention, the default input size in **torchgan.models** is a power of 2 and at least 16. Hence we shall be resizing the images to $3 \\times 64 \\times 64$.\n",
    "\n",
    "2. First, we will make a **center crop** of $160 \\times 160$ of the images. Next we resize the images to $3 \\times 64 \\times 64$. Direct resizing, affects the quality of the image. We want to generate faces, hence we want a major portion of the images to be filled by the faces.\n",
    "\n",
    "3. The output quality of GANs is improved when the images are normalized with a mean and standard deviation of **0.5**, thereby constraining pixel values of the input in the range $(-1, 1)$.\n",
    "\n",
    "Finally the **torchgan.trainer.Trainer** needs a **DataLoader** as input. So we are going to construct a DataLoader for the MNIST Dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import gzip\n",
    "import json\n",
    "import shutil\n",
    "import zipfile\n",
    "import argparse\n",
    "import requests\n",
    "import subprocess\n",
    "try:\n",
    "    from tqdm import tqdm\n",
    "except ImportError:\n",
    "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"tqdm\"])\n",
    "from six.moves import urllib\n",
    "\n",
    "\n",
    "def download_file_from_google_drive(id, destination):\n",
    "    URL = \"https://docs.google.com/uc?export=download\"\n",
    "    session = requests.Session()\n",
    "\n",
    "    response = session.get(URL, params={\"id\": id}, stream=True)\n",
    "    token = get_confirm_token(response)\n",
    "\n",
    "    if token:\n",
    "        params = {\"id\": id, \"confirm\": token}\n",
    "        response = session.get(URL, params=params, stream=True)\n",
    "\n",
    "    save_response_content(response, destination)\n",
    "\n",
    "\n",
    "def get_confirm_token(response):\n",
    "    for key, value in response.cookies.items():\n",
    "        if key.startswith(\"download_warning\"):\n",
    "            return value\n",
    "    return None\n",
    "\n",
    "\n",
    "def save_response_content(response, destination, chunk_size=32 * 1024):\n",
    "    total_size = int(response.headers.get(\"content-length\", 0))\n",
    "    with open(destination, \"wb\") as f:\n",
    "        for chunk in tqdm(\n",
    "            response.iter_content(chunk_size),\n",
    "            total=total_size,\n",
    "            unit=\"B\",\n",
    "            unit_scale=True,\n",
    "            desc=destination,\n",
    "        ):\n",
    "            if chunk:  # filter out keep-alive new chunks\n",
    "                f.write(chunk)\n",
    "\n",
    "\n",
    "def download_celeb_a(dirpath):\n",
    "    data_dir = \"celebA\"\n",
    "    if os.path.exists(os.path.join(dirpath, data_dir)):\n",
    "        print(\"Found Celeb-A - skip\")\n",
    "        return\n",
    "\n",
    "    filename, drive_id = \"img_align_celeba.zip\", \"0B7EVK8r0v71pZjFTYXZWM3FlRnM\"\n",
    "    save_path = os.path.join(dirpath, filename)\n",
    "\n",
    "    if os.path.exists(save_path):\n",
    "        print(\"[*] {} already exists\".format(save_path))\n",
    "    else:\n",
    "        download_file_from_google_drive(drive_id, save_path)\n",
    "\n",
    "    zip_dir = \"\"\n",
    "    with zipfile.ZipFile(save_path) as zf:\n",
    "        zip_dir = zf.namelist()[0]\n",
    "        zf.extractall(dirpath)\n",
    "    os.remove(save_path)\n",
    "    os.rename(os.path.join(dirpath, zip_dir), os.path.join(dirpath, data_dir))\n",
    "\n",
    "\n",
    "os.mkdir(\"./datasets\")",
    "\n",
    "download_celeb_a(\"./datasets\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = dsets.ImageFolder(\n",
    "    \"./datasets/celebA\",\n",
    "    transform=transforms.Compose(\n",
    "        [\n",
    "            transforms.CenterCrop(160),\n",
    "            transforms.Resize((64, 64)),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n",
    "        ]\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloader = data.DataLoader(dataset, batch_size=64, shuffle=True, num_workers=8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GENERATOR MODEL\n",
    "\n",
    "The SAGANGenerator receives an input noise vector of size $batch\\ size \\times encoding\\ dims$. The output must be a torch Tensor of size $batch\\ size \\times 3 \\times 64 \\times 64$ conforming to the input MNIST size. We transform the noise to the image in the following way:\n",
    "\n",
    "1. **Channel Dimension**: $encoding\\ dims \\rightarrow d \\rightarrow \\frac{d}{2} \\rightarrow \\frac{d}{4} \\rightarrow \\frac{d}{8} \\rightarrow 1$.\n",
    "2. **Image size**: $(1 \\times 1) \\rightarrow (4 \\times 4) \\rightarrow (8 \\times 8) \\rightarrow (16 \\times 16) \\rightarrow (32 \\times 32) \\rightarrow (64 \\times 64)$.\n",
    "\n",
    "We are going to use **LeakyReLU** as the activation. Using ReLU will kill most of the gradients and hence convergence is critically slowed down. Hence it is a valid choice for the activation. At the end of the model we use a **Tanh** activation. This ensures that the pixel values range from $(-1\\ to\\ 1)$.\n",
    "\n",
    "---\n",
    "\n",
    "**SelfAttention2d** is present out-of-box in torchgan. Self Attention simply learns an attention map to learn ot attend to certain places of the image. It contains 3 convolutional layers, named **query**, **key** and **value**. The weights of these layers are learned. We compute the attention map by:\n",
    "\n",
    "$$attention = softmax((query(x)^T \\times key(x)))$$\n",
    "$$output = \\gamma \\times value(x) \\times attention + x $$\n",
    "\n",
    "$\\gamma$ is simply a scaling parameter which is also learning. Its value is initially set to 0. As result in the beginning, the Self Attention Layer does nothing. As the network starts learning the features the value of $\\gamma$ increases and so does the effect of the attention map. This is done to ensure that the model initially focusses on learning the local features (which are relatively easy to learn). Once it has learnt to detect the local features it focusses on learning the more\n",
    "global features.\n",
    "\n",
    "For more information on SelfAttention2d, use `help(SelfAttention2d)`.\n",
    "\n",
    "---\n",
    "\n",
    "The Generator model needs 2 inputs, the $1^{st}$ one is the **encoding_dims** and the $2^{nd}$ one is the **label_type**. Since SAGAN is not a conditional model, we are going to set the **label_type** as **'none'**. For conditioned models there are 2 options - **'generated'** and **'required'**. We shall discuss them in detail in later tutorials.\n",
    "\n",
    "The training of SAGAN is same as the standard GANs, it is only an architectural improvement. Hence subclassing the Generator, allows us to use the default **train_ops**. We shall be discussing on how to write a custom **train_ops** in a future tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SAGANGenerator(Generator):\n",
    "    def __init__(self, encoding_dims=100, step_channels=64):\n",
    "        super(SAGANGenerator, self).__init__(encoding_dims, \"none\")\n",
    "        d = int(step_channels * 8)\n",
    "        self.model = nn.Sequential(\n",
    "            SpectralNorm2d(\n",
    "                nn.ConvTranspose2d(self.encoding_dims, d, 4, 1, 0, bias=True)\n",
    "            ),\n",
    "            nn.BatchNorm2d(d),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            SpectralNorm2d(nn.ConvTranspose2d(d, d // 2, 4, 2, 1, bias=True)),\n",
    "            nn.BatchNorm2d(d // 2),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            SpectralNorm2d(nn.ConvTranspose2d(d // 2, d // 4, 4, 2, 1, bias=True)),\n",
    "            nn.BatchNorm2d(d // 4),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            SelfAttention2d(d // 4),\n",
    "            SpectralNorm2d(nn.ConvTranspose2d(d // 4, d // 8, 4, 2, 1, bias=True)),\n",
    "            nn.BatchNorm2d(d // 8),\n",
    "            SelfAttention2d(d // 8),\n",
    "            SpectralNorm2d(nn.ConvTranspose2d(d // 8, 3, 4, 2, 1, bias=True)),\n",
    "            nn.Tanh(),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.view(-1, x.size(1), 1, 1)\n",
    "        return self.model(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DISCRIMINATOR MODEL\n",
    "\n",
    "The SAGANDiscriminator receives an image in the form of a torch Tensor of size $batch\\ size \\times 3 \\times 64 \\times 64$. It must output a tensor of size $batch\\ size$, where each value corresponds to the confidence of whether or not that is image is real. Since we use a **Tanh** activation at the end of the model, the outputs **cannot be interpreted as probabilities**. Instead we interpret them as follows:\n",
    "\n",
    "1. A **Higher** value (*closer to 1.0*) indicates that the discriminator believes the image to be a real one.\n",
    "2. A **Lower** value (*closer to -1.0*) indicates that the discriminator believes the image to be a fake one.\n",
    "\n",
    "---\n",
    "\n",
    "For reasons same as above we use a **LeakyReLU** activation. The conversion of the image tensor to the confidence scores are as follows:\n",
    "\n",
    "1. **Channel Dimension**: $1 \\rightarrow d \\rightarrow 2 \\times d \\rightarrow 4 \\times d \\rightarrow 8 \\times d \\rightarrow 1$.\n",
    "2. **Image size**: $(64 \\times 64) \\rightarrow (32 \\times 32) \\rightarrow (16 \\times 16) \\rightarrow (8 \\times 8) \\rightarrow (4 \\times 4) \\rightarrow (1 \\times 1)$.\n",
    "\n",
    "---\n",
    "\n",
    "The Discriminator also needs 2 inputs. The $1^{st}$ one is the channel dimension of the input image. In our case, we have grayscale images so this value is *1*. The $2^{nd}$ argument is the **label_type**, which is again **'none'**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SAGANDiscriminator(Discriminator):\n",
    "    def __init__(self, step_channels=64):\n",
    "        super(SAGANDiscriminator, self).__init__(3, \"none\")\n",
    "        d = step_channels\n",
    "        self.model = nn.Sequential(\n",
    "            SpectralNorm2d(nn.Conv2d(self.input_dims, d, 4, 2, 1, bias=True)),\n",
    "            nn.BatchNorm2d(d),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            SpectralNorm2d(nn.Conv2d(d, d * 2, 4, 2, 1, bias=True)),\n",
    "            nn.BatchNorm2d(d * 2),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            SpectralNorm2d(nn.Conv2d(d * 2, d * 4, 4, 2, 1, bias=True)),\n",
    "            nn.BatchNorm2d(d * 4),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            SelfAttention2d(d * 4),\n",
    "            SpectralNorm2d(nn.Conv2d(d * 4, d * 8, 4, 2, 1, bias=True)),\n",
    "            nn.BatchNorm2d(d * 8),\n",
    "            SelfAttention2d(d * 8),\n",
    "            SpectralNorm2d(nn.Conv2d(d * 8, 1, 4, 1, 0, bias=True)),\n",
    "            nn.Tanh(),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.model(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MODEL PARAMETERS & HYPERPARAMETERS\n",
    "\n",
    "In this section we are going to define how the model is going to look like, what their optimizers are going to be and every other possible hyperparameter. We use a dictionary to feed the input into the trainer. We believe it to be a much cleaner approach than individual arguments which tend to limit the capacity of the Trainer.\n",
    "\n",
    "---\n",
    "\n",
    "We follow a standard naming system for parsing arguments. **name** refers to the **class** whose object is to be created. **args** contains the arguments that will be fed into the **class** while instantiating it. The keys refer to the **variable names** with which the **object** shall be stored in the **Trainer**. This storage pattern allows us to have customs **train_ops** for Losses with no furthur headaches, but thats again for a future discussion. The items in the dictionary are allowed to have the following keys:\n",
    "\n",
    "1. \"name\": The class name for the model.\n",
    "2. \"args\": Arguments fed into the class.\n",
    "3. \"optimizer\": A dictionary containing the following key-value pairs\n",
    "    * \"name\"\n",
    "    * \"args\"\n",
    "    * \"var\": Variable name for the optimizer. This is an optional argument. If this is not provided, we assign the optimizer the name **optimizer_{}** where {} refers to the variable name of the model.\n",
    "    * \"scheduler\": Optional scheduler associated with the optimizer. Again this is a dictionary with the following keys\n",
    "        * \"name\"\n",
    "        * \"args\"\n",
    "\n",
    "---\n",
    "\n",
    "Lets see the interpretation of the `network_params` defined in the following code block. This dictionary gets interpreted as follows in the Trainer:\n",
    "\n",
    "```python\n",
    "generator = SAGANGenerator(step_channels = 32)\n",
    "optimizer_generator = Adam(generator.parameters(), lr = 0.0001, betas = (0.0, 0.999))\n",
    "\n",
    "discriminator = SAGANDiscriminator(step_channels = 32)\n",
    "optimizer_discriminator = Adam(discriminator.parameters(), lr = 0.0004, betas = (0.0, 0.999))\n",
    "```\n",
    "\n",
    "As observed we are using the **TTUR (Two Timescale Update  Rule)** for training the models. Hence we can easily customize most of the training pipeline even with such high levels of abstraction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "network_params = {\n",
    "    \"generator\": {\n",
    "        \"name\": SAGANGenerator,\n",
    "        \"args\": {\"step_channels\": 32},\n",
    "        \"optimizer\": {\"name\": Adam, \"args\": {\"lr\": 0.0001, \"betas\": (0.0, 0.999)}},\n",
    "    },\n",
    "    \"discriminator\": {\n",
    "        \"name\": SAGANDiscriminator,\n",
    "        \"args\": {\"step_channels\": 32},\n",
    "        \"optimizer\": {\"name\": Adam, \"args\": {\"lr\": 0.0004, \"betas\": (0.0, 0.999)}},\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda:0\")\n",
    "    # Use deterministic cudnn algorithms\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    epochs = 20\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "    epochs = 10\n",
    "\n",
    "print(\"Device: {}\".format(device))\n",
    "print(\"Epochs: {}\".format(epochs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to feed a list of losses to the trainer. These losses control the optimization of the model weights. We are going to use the following losses, all of which are defined in the **torchgan.losses** module.\n",
    "\n",
    "1. WassersteinGeneratorLoss\n",
    "2. WassersteinDiscriminatorLoss\n",
    "\n",
    "We need the **clip** parameter to enforce the **Lipschitz condition**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "losses_list = [\n",
    "    WassersteinGeneratorLoss(),\n",
    "    WassersteinDiscriminatorLoss(clip=(-0.01, 0.01)),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## VISUALIZE THE TRAINING DATA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot some of the training images\n",
    "real_batch = next(iter(dataloader))\n",
    "plt.figure(figsize=(8, 8))\n",
    "plt.axis(\"off\")\n",
    "plt.title(\"Training Images\")\n",
    "plt.imshow(\n",
    "    np.transpose(\n",
    "        vutils.make_grid(\n",
    "            real_batch[0].to(device)[:64], padding=2, normalize=True\n",
    "        ).cpu(),\n",
    "        (1, 2, 0),\n",
    "    )\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TRAINING THE GENERATOR & DISCRIMINATOR\n",
    "\n",
    "Next we simply feed the network descriptors and the losses we defined previously into the Trainer. Then we pass the **CelebA DataLoader** to the trainer object and wait for training to complete.\n",
    "\n",
    "---\n",
    "\n",
    "Important information for visualizing the performance of the GAN will be printed to the console. The best and recommended way to visualize the training is to use **tensorboardX**. It plots all the data and periodically displays the generated images. It allows us to track failure of the model early."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    network_params, losses_list, sample_size=64, epochs=epochs, device=device\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer(dataloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.complete()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## VISUALIZE THE GENERATED DATA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Grab a batch of real images from the dataloader\n",
    "real_batch = next(iter(dataloader))\n",
    "\n",
    "# Plot the real images\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.axis(\"off\")\n",
    "plt.title(\"Real Images\")\n",
    "plt.imshow(\n",
    "    np.transpose(\n",
    "        vutils.make_grid(\n",
    "            real_batch[0].to(device)[:64], padding=5, normalize=True\n",
    "        ).cpu(),\n",
    "        (1, 2, 0),\n",
    "    )\n",
    ")\n",
    "\n",
    "# Plot the fake images from the last epoch\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.axis(\"off\")\n",
    "plt.title(\"Fake Images\")\n",
    "plt.imshow(plt.imread(\"{}/epoch{}_generator.png\".format(trainer.recon, trainer.epochs)))\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
